package org.wyz.spark;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.broadcast.Broadcast;
import scala.Tuple2;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;


/**
 * 大部分代码可以复用Join的
 */
public class JoinWithBroadcast {
    private static final String[] ARGS = {"data/input/join/user.txt", "data/input/join/user-exam.txt", "data/output/join-with-broadcast"};

    // id拿走，剩下的字段构成一个List
    public static List<String> withoutId(int idIndex, String[] ss) {
        ArrayList<String> res = new ArrayList<>();
        for (int i = 0; i < ss.length; i++) {
            if (i == idIndex) continue;
            res.add(ss[i]);
        }
        return res;
    }

    // 2个List合并为1个
    public static List<String> mergeList(List<String> a, List<String> b) {
        ArrayList<String> res = new ArrayList<>();
        res.addAll(a);
        res.addAll(b);
        return res;
    }


    public static void main(String[] args) {
        SparkConf conf = new SparkConf().setMaster("local").setAppName("wyz-join-2");
        JavaSparkContext sc = new JavaSparkContext(conf);


        // 小表
        // collectAsMap之后计算结果会被返回给Driver
        Map<String, List<String>> map = sc.textFile(ARGS[0]).mapToPair(a -> {
            String[] ss = a.split("\\s");
            return new Tuple2<>(ss[0], withoutId(0, ss));
        }).collectAsMap();


        // 广播小表
        // 由Driver执行
        Broadcast<Map<String, List<String>>> broadcast = sc.broadcast(map);
        System.out.println("=" + System.identityHashCode(broadcast));


        // 大表
        JavaPairRDD<String, List<String>> rdd0 = sc.textFile(ARGS[1]).mapToPair(a -> {
            String[] ss = a.split("\\s");
            return new Tuple2<>(ss[0], withoutId(0, ss));
        });


        // 不能用join算子了，因为join算子会shuffle，所以改用map
        JavaPairRDD<String, List<String>> res = rdd0.mapToPair(a -> {
            Map<String, List<String>> m = broadcast.value();
            System.out.println("===" + System.identityHashCode(broadcast));
            if (m.containsKey(a._1)) {
                return new Tuple2<>(a._1, mergeList(a._2, m.get(a._1)));
            }
            return new Tuple2<>(a._1, null);
        }).filter(a -> a._2 != null);

        res.map(a -> a._1 + " " + a._2).saveAsTextFile(ARGS[2]);


        sc.stop();
    }
}
